'''
Created on Aug 4, 2014

@author: Siyuan Qi
'''

import os
import time
import copy
import shutil

# Dependencies
import numpy as np
import matplotlib.pyplot as plt

# Local imports
import proposal
import grammarModel as sg
import visualizeScene as vs
import likelihoodModel as lm


class MCMCParser():
	def __init__(self, xml, img, workPath):
		self.saveTempResult = False
		
		self.xml = xml
		self.imgName = img.rstrip()
		self.directoryWork = workPath
		self.directoryResult = workPath+"results/"+os.path.splitext(self.imgName)[0]+"/"
		self.directoryTemp = self.directoryResult+"tempResults/"
		
		# Create the results folders if it does not exist
		if not os.path.exists(self.directoryResult):
			os.makedirs(self.directoryResult)
		if self.saveTempResult:
			if not os.path.exists(self.directoryTemp):
				os.makedirs(self.directoryTemp)
			
		# Copy the original image to the results folder
		shutil.copyfile(workPath+"images/"+self.imgName, self.directoryResult+self.imgName)
		
		self.outputFile = open(self.directoryResult+"output.txt", 'w')
		
		# Load all the data in generated proposals(background proposals and cuboid proposals)
		self.data = proposal.loadData(self.directoryWork+'output/Images/', self.imgName)
		
		# Initialize the temperature and the energy list
		self.T = 10.
		self.energies = []
		
		# Initialize the scene model and the likelihood model
		self.sceneModel = sg.GrammarModel(self.xml, self.data[4], self.saveTempResult, self.directoryTemp)
		self.likeModel = lm.LikelihoodModel(self.data[0:4], self.directoryWork, self.imgName)
		
		# Need to build distributions for the proposals
		self.bgDist = self.data[5]
		self.cubDist = self.data[6]
		self.recProps = self.data[7]
		
		# Fit non-parametric distributions of background and cuboid proposals
		self.bgDist.fitDist()
		self.cubDist.addObjectFilters(self.xml)
		self.cubDist.fitDist(0.1)
		
		# Initialize(pre-sample) the parse tree
		self.updateTemperature()
		self.sceneModel.initializeParseTree(self.bgDist, self.cubDist)
		self.newSceneModel = None
		
		self.newPriorEnergy = -self.sceneModel.calLogPrior()
		self.newLikelihoodEnergy = self.likeModel.calLikelihoodEnergy(self.sceneModel.getEntities())
		self.newPhysicsEnergy = self.sceneModel.calPhysicsEnergy()
		self.newEnergy = self.newPriorEnergy + self.newLikelihoodEnergy + self.newPhysicsEnergy
		self.oldEnergy = self.newEnergy
		
		self.minEnergy = np.inf
		self.conRejections = 0
		
		if self.saveTempResult: 
			self.cubDist.saveFrames(self.directoryTemp, self.likeModel, 100, 10)
		
	def updateTemperature(self):
		# Update the cuboid distribution and the position distributions in 
		# the grammar model if the temperature is larger than 1. Otherwise
		# the temperature only affects the acceptance ratio.
		if self.T >= 1:
			self.cubDist.fitDist(0.1*self.T)
			self.sceneModel.updateTemperature(self.T, self.cubDist)

	def cal_new_energy(self):
		self.newPriorEnergy = -self.newSceneModel.calLogPrior()
		self.newLikelihoodEnergy = self.likeModel.calLikelihoodEnergy(self.newSceneModel.getEntities())
		self.newPhysicsEnergy = self.newSceneModel.calPhysicsEnergy()
		self.newEnergy = self.newPriorEnergy + self.newLikelihoodEnergy + self.newPhysicsEnergy
		
	def calAcceptanceRate(self):
		self.cal_new_energy()
		posteriorRatio = -(self.newEnergy - self.oldEnergy)/self.T
		p = posteriorRatio+self.proposeProbLogRatio
		info = "priorEnergy, likelihoodEnergy, physicsEnergy, proposeProbLogRatio, p: {:f}, {:f}, {:f}, {:f}, {:f}\n".format(self.newPriorEnergy, self.newLikelihoodEnergy, self.newPhysicsEnergy, self.proposeProbLogRatio, p)
		if p>0:
			return 1, info
		else:
			return np.exp(p), info
		
	def propose(self, iteration):
		self.newSceneModel = copy.deepcopy(self.sceneModel)
		self.proposeProbLogRatio, method = self.newSceneModel.propose(iteration)
		
		rate, info = self.calAcceptanceRate()
		
		if np.random.rand() <= rate:
			self.sceneModel = self.newSceneModel
			self.oldEnergy = self.newEnergy
			self.conRejections = 0
			accepted = True
		else:
			self.conRejections += 1
			accepted = False
		
		# ====================== Write outputs ======================
		# Calculate the new energy after acceptance/rejection
		self.energies.append(self.oldEnergy)
		# Save this result if it hits the lowest energy so far
		if self.newEnergy < self.minEnergy:
			self.minEnergy = self.newEnergy
			self.saveResult("lowestEnergy")
		
		# Write down information for this iteration into output file
		self.outputFile.write("\nIteration: {:d}\n".format(iteration))
		self.outputFile.write(method+"\n")
		self.outputFile.write(info)
		
		if accepted:
			self.outputFile.write("Accepted\n")
		else:
			self.outputFile.write("Rejected\n")
		
		return True
		
	def run(self):
		t0 = time.time()
		self.outputFile.write("Start: "+self.imgName+"\n")
		self.saveResult("original")
		
		i = 0
		while i<=20: #1999:
			# print i
			if self.propose(i):
				i += 1
				# self.T -= 0.005
				# self.T -= 0.02
				self.T -= 1
			
			if i%25 == 0:
				self.updateTemperature() 
			
			if self.saveTempResult:
				self.saveTemperaryResult(i)
		
		self.T = 0.001
		self.updateTemperature() 
		self.conRejections = 0
		while True:
			# print i
			# if self.conRejections >= 300 or i >= 3000:
			if self.conRejections >= 3 or i >= 30:
				break
			
			if self.propose(i):
				i += 1
			
			if self.saveTempResult:
				self.saveTemperaryResult(i)
			
		t1 = time.time()

		self.energy_descent()
		self.cuboid_proposal_descent()

		self.saveResult("final", True)
		t2 = time.time()
		self.outputFile.write("\nParsing time: {:f}\n".format(t1 - t0))
		self.outputFile.write("Total time: {:f}\n".format(t2 - t0))
		
	
	def saveTemperaryResult(self, iteration):
		self.likeModel.saveTemp(self.directoryTemp, iteration)
		
		plt.plot(self.energies)
		plt.savefig(self.directoryTemp+"energies_"+str(iteration)+".png")
		plt.clf()

	def saveResult(self, name = None, saveAll = False, show = False):
		if saveAll:
			# Write energies to file
			energyFile = open(self.directoryResult+"energies.txt", "w")
			for e in self.energies:
				energyFile.write("{}\n".format(e))
			
			self.sceneModel.drawGraph('finalGraph.png')
			self.sceneModel.writeEntitiesToFile(self.directoryResult+"sampledCuboids.txt")
			plt.plot(self.energies)
			plt.savefig(self.directoryResult+"energies.png")
			plt.clf()
			vs.savePlotToFile(self.sceneModel.getEntities(), self.directoryResult+"sampledRoom.png")
			
			self.likeModel.calLikelihoodEnergy(self.sceneModel.getEntities())
			self.likeModel.saveMap(self.directoryResult, name, True)
		else:
			self.likeModel.saveMap(self.directoryResult, name)
		
		if show:
			plt.show()
		
	def visualize(self):
		vs.plotFromList(self.sceneModel.getEntities())


	def energy_descent(self):
		# TODO: implement this method
		step_size = 0.02
		pt_entities = self.sceneModel.get_entities()

		for entity in pt_entities:
			print entity.node.symbol
		pass


	def cuboid_proposal_descent(self):
		# TODO: implement this method
		pass



